from zipline.api import symbol, order_target_percent, record, get_datetime, schedule_function, date_rules
from zipline import run_algorithm
import pandas as pd
import numpy as np

def initialize(context):
    # 策略参数
    context.lookback_years = 10
    context.asset = symbol('SPY')  # 示例资产
    
    # 每月初执行再平衡
    schedule_function(rebalance, date_rules.month_start())

def rebalance(context, data):
    # 获取历史价格数据
    prices = data.history(
        context.asset, 
        'price', 
        context.lookback_years * 252,  # 每年约252个交易日
        '1d'
    )
    
    # 计算月度收益
    monthly_prices = prices.resample('M').last()
    monthly_returns = monthly_prices.pct_change().dropna()
    
    # 计算历史各月平均收益
    seasonal_pattern = monthly_returns.groupby(
        monthly_returns.index.month
    ).mean().sort_values()
    
    # 确定最佳和最差月份
    best_month = seasonal_pattern.idxmax()
    worst_month = seasonal_pattern.idxmin()
    
    # 当前月份
    current_month = get_datetime().month
    
    # 交易逻辑
    if current_month == best_month:
        order_target_percent(context.asset, 1.0)  # 做多
    elif current_month == worst_month:
        order_target_percent(context.asset, -1.0)  # 做空
    else:
        order_target_percent(context.asset, 0.0)  # 空仓

# 示例回测配置
if __name__ == '__main__':
    start = pd.Timestamp('2010-01-01', tz='utc')
    end = pd.Timestamp('2023-01-01', tz='utc')
    
    results = run_algorithm(
        start=start,
        end=end,
        initialize=initialize,
        capital_base=10000,
        data_frequency='daily',
        bundle='quandl'
    )